
#include "global_avg_pool_tiling.h"
#include "register/op_def_registry.h"
#include "tiling_common.h"
namespace optiling
{
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{

    GlobalAvgPoolTilingData tiling;
    context->SetBlockDim(1);
    auto x1_shape = context->GetInputShape(0)->GetStorageShape();
    auto n = x1_shape.GetDim(0);
    auto c = x1_shape.GetDim(1);
    auto d = x1_shape.GetShapeSize() / n / c;
    tiling.set_nxc(n * c);
    tiling.set_d(d);
    
    COMMON_TILING_GET_PLATFORM;
    COMMON_TILING_GET_UB_SIZE;
    auto dtype = context->GetInputDesc(0)->GetDataType();
    auto dtypeSize = GetSizeByDataType(dtype);
    int64_t alignSize = 256 / dtypeSize;
    // xque(2) + worklocal(2/alignSize)
    double tileLenthCount = 2.0 + (2 / double(alignSize));
    if (dtype == ge::DataType::DT_FLOAT16)
    {
        // xque(2) + worklocal(2/alignSize) * 2 + casted_x(2)
        tileLenthCount = 2.0 + (2 / double(alignSize)) * 2 + 2.0;
    }
    // worklocal对齐(256) + sum_value(32 * 2)
    auto extra = 256 + 32 * 2 + 32 * 2;
    auto tileLength = int64_t(double(ub_size - extra) / tileLenthCount) / 256 * 256 / dtypeSize;
    if (d < tileLength)
    {
        tileLength = ALIGN_TO(d, alignSize);
    }
    tiling.set_tileLength(tileLength);
    auto finalLength = d % tileLength;
    if (finalLength == 0)
    {
        finalLength = d;
    }
    tiling.set_finalLength(finalLength);

    tiling_zero_workspace(context);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    std::cout << "nxc = " << tiling.get_nxc() << std::endl;
    std::cout << "d = " << tiling.get_d() << std::endl;
    std::cout << "tileLength = " << tiling.get_tileLength() << std::endl;
    std::cout << "finalLength = " << tiling.get_finalLength() << std::endl;

    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge
{
static ge::graphStatus InferShape(gert::InferShapeContext *context)
{
    const gert::Shape *x1_shape = context->GetInputShape(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops
{
class GlobalAvgPool : public OpDef
{
  public:
    explicit GlobalAvgPool(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
    }
};

OP_ADD(GlobalAvgPool);
} // namespace ops
